from typing import Callable, Union
import numpy as np
from tianshou.data import Batch, to_numpy


class RulePolicyBase:
    def set_ril_pt(self, pt):
        """记得在 train_fn 里 set pt！"""
        self.ril_pt = pt

    def __init__(
            self,
            rule_fn: Callable = None
    ):
        """
        :param rule_fn: ril 的函数，输入 state, action 返回 action, mask, explorable
        """
        self.ril_pt = 0.0  # ril_pt 是在 train_fn 里 set 的
        self.rule_fn = rule_fn

    def process_rule(
            self,
            batch: Batch,
            act: np.ndarray
    ) -> (np.ndarray, Union[np.ndarray, None]):
        """
        :return: (act: 修改后的动作, mask: 屏蔽的动作, explorable: 是否允许随机动作)
        每个 mask[i] 是形如 [True, True, False, ...] 的 ndarray，True 代表这个动作可以选，反之不可
        """
        obs_ = batch.obs.obs if hasattr(batch.obs, "obs") else batch.obs
        if hasattr(batch, "info") and hasattr(batch.info, "origin"):
            obs_ = batch.info.origin
        mask = [None] * len(act)  # 禁用某些动作
        explorable = [True] * len(act)  # 是否允许探索

        if self.rule_fn is not None and np.random.rand() < self.ril_pt:
            for i in range(len(act)):
                act[i], mask[i], explorable[i] = self.rule_fn(obs_[i], act[i])
        return act, np.array(mask), explorable
